Skip to content

LCORE-2309: Added Pydantic AI Bridge#1817

Merged
tisnik merged 4 commits into
lightspeed-core:mainfrom
jrobertboos:lcore-2309
Jun 4, 2026
Merged

LCORE-2309: Added Pydantic AI Bridge#1817
tisnik merged 4 commits into
lightspeed-core:mainfrom
jrobertboos:lcore-2309

Conversation

@jrobertboos
Copy link
Copy Markdown
Contributor

@jrobertboos jrobertboos commented May 29, 2026

Description

Type of change

  • Refactor
  • New feature
  • Bug fix
  • CVE fix
  • Optimization
  • Documentation Update
  • Configuration Update
  • Bump-up service version
  • Bump-up dependent library
  • Bump-up library or tool used for development (does not change the final image)
  • CI configuration change
  • Konflux configuration change
  • Unit tests improvement
  • Integration tests improvement
  • End to end tests improvement
  • Benchmarks improvement

Tools used to create PR

Identify any AI code assistants used in this PR (for transparency and review context)

  • Assisted-by: (e.g., Claude, CodeRabbit, Ollama, etc., N/A if not used)
  • Generated by: (e.g., tool name and version; N/A if not used)

Related Tickets & Documents

  • Related Issue #
  • Closes #

Checklist before requesting a review

  • I have performed a self-review of my code.
  • PR has passed all pre-merge test jobs.
  • If it is a core feature, I have added thorough tests.

Testing

  • Please provide detailed steps to perform tests related to this code change.
  • How were the fix/results from this change verified? Please provide relevant screenshots or results.

Summary by CodeRabbit

  • New Features
    • Added support for creating Pydantic AI agents configured for Llama Stack environments, with automatic parameter mapping for API compatibility.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 29, 2026

Review Change Stack

Walkthrough

A new module src/utils/pydantic_ai.py is added containing Pydantic AI integration utilities for Llama Stack. The module exports build_agent to construct agents from HTTP or library clients and Responses API parameters. Supporting functions include provider factory, settings mapper, and a filtered extra fields constant. Comprehensive unit tests verify all mappings and client handling paths.

Changes

Pydantic AI Llama Stack Integration Utilities

Layer / File(s) Summary
Pydantic AI integration utilities
src/utils/pydantic_ai.py
_LLS_RESPONSES_EXTRA_FIELDS defines an allowlist of Llama Stack field names for extra_body injection. _llama_stack_provider_from_client derives a LlamaStackProvider from either an AsyncLlamaStackClient (HTTP server mode with base URL /v1 normalization and API key fallback) or AsyncLlamaStackAsLibraryClient (in-process library mode, reusing the client's HTTP client). _model_settings_from_responses_params maps ResponsesApiParams to OpenAIResponsesModelSettings, conditionally translating field names (max_output_tokensmax_tokens, storeopenai_store, previous_response_idopenai_previous_response_id), converting extra_headers to dict, and injecting filtered extra_body when applicable. build_agent wires the provider and mapped settings into an OpenAIResponsesModel and returns an Agent configured with instructions and defer_model_check=True.
Integration utilities tests
tests/unit/utils/test_pydantic_ai.py
Test imports and fixtures establish context for testing internal symbols. Tests for _llama_stack_provider_from_client verify HTTP client forwarding, base URL content, API key defaulting to "not-needed" when absent, and library-client passthrough. Tests for _model_settings_from_responses_params verify parameter mapping, openai_store derivation, conditional extra_body exclusion when empty, field renames, and allowlist-filtered field injection. Tests for _LLS_RESPONSES_EXTRA_FIELDS assert it is a frozenset matching an exact expected key set. Tests for build_agent verify agent construction for both remote and library client modes and instruction incorporation.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

Suggested reviewers

  • asimurka
  • tisnik
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title 'LCORE-2309: Added Pydantic AI Bridge' directly and clearly describes the main change—addition of a Pydantic AI integration bridge—which is confirmed by the new module and tests.
Docstring Coverage ✅ Passed Docstring coverage is 93.06% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
✨ Simplify code
  • Create PR with simplified code

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Comment thread src/pydantic_ai_lightspeed/llamastack/_provider.py Fixed
Comment thread src/pydantic_ai_lightspeed/llamastack/_provider.py Fixed
Comment thread src/pydantic_ai_lightspeed/llamastack/_provider.py Fixed
Comment thread src/pydantic_ai_lightspeed/llamastack/_provider.py Fixed
Comment thread src/pydantic_ai_lightspeed/llamastack/_transport.py Fixed
Comment thread src/pydantic_ai_lightspeed/llamastack/_transport.py Fixed
@jrobertboos jrobertboos force-pushed the lcore-2309 branch 2 times, most recently from 402b0f2 to 53e1b86 Compare June 1, 2026 15:25
@jrobertboos jrobertboos marked this pull request as ready for review June 1, 2026 15:55
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/pydantic_ai_lightspeed/llamastack/_transport.py`:
- Around line 89-94: The header presence check treats names case-sensitively and
only checks two spellings, so normalize header names before deciding to inject
provider data: in the block that references self._client.provider_data and the
headers dict, compute a lowercase-key view (e.g., {k.lower(): v for k, v in
headers.items()} or use any(k.lower() == "x-llamastack-provider-data" for k in
headers)) and check for "x-llamastack-provider-data" in that normalized view;
only if absent, add the provider data using the canonical header
"X-LlamaStack-Provider-Data" to headers. Ensure you reference headers and
self._client.provider_data when applying the change.
- Around line 22-40: The _AsyncByteStream currently yields from the wrapped
async generator but doesn't forward closure; implement an async aclose(self) on
class _AsyncByteStream that calls and awaits self._gen.aclose() if the wrapped
generator has an aclose attribute to ensure httpx.Response.aclose() triggers
generator cleanup. In LlamaStackLibraryTransport.handle_async_request change the
provider-data header existence check to be case-insensitive (e.g., inspect
request.headers keys lowercased or use any(k.lower() ==
"x-llamastack-provider-data" for k in request.headers)) before injecting
"X-LlamaStack-Provider-Data" so you don't add a duplicate header when callers
use different casing.

In `@src/utils/pydantic_ai.py`:
- Around line 41-45: The helper _llama_stack_provider_from_client currently
reaches into AsyncLlamaStackClient._client (protected) when building a
LlamaStackProvider; change it to rely on the public HTTP-injection path instead:
update _llama_stack_provider_from_client to accept a public http_client
parameter (or obtain the client's configured http_client via its public API),
and pass that http_client into LlamaStackProvider(base_url=..., api_key=...,
http_client=...) rather than using client._client; ensure callers of
_llama_stack_provider_from_client are updated to provide the public http_client
argument.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 2d1bdbfc-90ec-412a-b14c-474e9f95a843

📥 Commits

Reviewing files that changed from the base of the PR and between 4410134 and 773aca1.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (11)
  • pyproject.toml
  • src/pydantic_ai_lightspeed/__init__.py
  • src/pydantic_ai_lightspeed/llamastack/__init__.py
  • src/pydantic_ai_lightspeed/llamastack/_provider.py
  • src/pydantic_ai_lightspeed/llamastack/_transport.py
  • src/utils/pydantic_ai.py
  • tests/unit/pydantic_ai_lightspeed/__init__.py
  • tests/unit/pydantic_ai_lightspeed/llamastack/__init__.py
  • tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py
  • tests/unit/pydantic_ai_lightspeed/llamastack/test_transport.py
  • tests/unit/utils/test_pydantic_ai.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (12)
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
  • GitHub Check: unit_tests (3.13)
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
  • GitHub Check: build-pr
  • GitHub Check: list_outdated_dependencies
  • GitHub Check: spectral
  • GitHub Check: E2E: server mode / ci / group 2
  • GitHub Check: E2E: library mode / ci / group 1
  • GitHub Check: E2E: server mode / ci / group 3
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: E2E: library mode / ci / group 3
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Use async def for I/O operations and external API calls
Use standard log levels with clear purposes: debug() for diagnostic info, info() for program execution, warning() for unexpected events, error() for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes: Configuration, Error/Exception, Resolver, Interface
Abstract classes must use ABC with @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/pydantic_ai_lightspeed/llamastack/__init__.py
  • src/pydantic_ai_lightspeed/__init__.py
  • src/utils/pydantic_ai.py
  • src/pydantic_ai_lightspeed/llamastack/_provider.py
  • src/pydantic_ai_lightspeed/llamastack/_transport.py
src/**/__init__.py

📄 CodeRabbit inference engine (AGENTS.md)

Package __init__.py files must contain brief package descriptions

Files:

  • src/pydantic_ai_lightspeed/llamastack/__init__.py
  • src/pydantic_ai_lightspeed/__init__.py
tests/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Use pytest.mark.asyncio marker for async tests

Files:

  • tests/unit/pydantic_ai_lightspeed/llamastack/__init__.py
  • tests/unit/pydantic_ai_lightspeed/__init__.py
  • tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py
  • tests/unit/pydantic_ai_lightspeed/llamastack/test_transport.py
  • tests/unit/utils/test_pydantic_ai.py
🔇 Additional comments (5)
pyproject.toml (1)

82-83: LGTM!

src/pydantic_ai_lightspeed/__init__.py (1)

1-1: LGTM!

src/pydantic_ai_lightspeed/llamastack/__init__.py (1)

1-6: LGTM!

tests/unit/pydantic_ai_lightspeed/__init__.py (1)

1-1: LGTM!

tests/unit/pydantic_ai_lightspeed/llamastack/__init__.py (1)

1-1: LGTM!

Comment thread src/pydantic_ai_lightspeed/llamastack/_transport.py
Comment thread src/pydantic_ai_lightspeed/llamastack/_transport.py
Comment thread src/utils/pydantic_ai.py
Comment on lines +86 to +87
lib_http_client = httpx.AsyncClient(
transport=transport,
Copy link
Copy Markdown
Contributor

@asimurka asimurka left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/utils/pydantic_ai.py`:
- Line 35: The single-line docstrings for _llama_stack_provider_from_client and
_model_settings_from_responses_params must be expanded to follow Google Python
docstring conventions: update each function to include Parameters (documenting
args such as client, responses_params, etc. with types and purpose), Returns
(describe the return type and what it represents), and Raises (list any
exceptions the function may raise); ensure the wording matches existing
parameter names and types used in those functions and keep the descriptions
concise and accurate so the new docstrings fully replace the current one-line
descriptions.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 31a18f2c-2ccf-4a87-951f-c7039dd5df7d

📥 Commits

Reviewing files that changed from the base of the PR and between 773aca1 and 08bee6f.

📒 Files selected for processing (2)
  • src/utils/pydantic_ai.py
  • tests/unit/utils/test_pydantic_ai.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (11)
  • GitHub Check: shellcheck
  • GitHub Check: unit_tests (3.13)
  • GitHub Check: mypy
  • GitHub Check: ruff
  • GitHub Check: Pylinter
  • GitHub Check: integration_tests (3.12)
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: E2E: library mode / ci / group 3
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
🧰 Additional context used
📓 Path-based instructions (2)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Use async def for I/O operations and external API calls
Use standard log levels with clear purposes: debug() for diagnostic info, info() for program execution, warning() for unexpected events, error() for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes: Configuration, Error/Exception, Resolver, Interface
Abstract classes must use ABC with @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/utils/pydantic_ai.py
tests/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Use pytest.mark.asyncio marker for async tests

Files:

  • tests/unit/utils/test_pydantic_ai.py
🧠 Learnings (4)
📓 Common learnings
Learnt from: CR
Repo: lightspeed-core/lightspeed-stack PR: 0
File: AGENTS.md:0-0
Timestamp: 2026-05-06T06:57:52.173Z
Learning: Applies to src/**/*.py : Llama Stack imports: Use `from llama_stack_client import AsyncLlamaStackClient`
📚 Learning: 2026-05-06T06:57:52.173Z
Learnt from: CR
Repo: lightspeed-core/lightspeed-stack PR: 0
File: AGENTS.md:0-0
Timestamp: 2026-05-06T06:57:52.173Z
Learning: Applies to src/**/*.py : Llama Stack imports: Use `from llama_stack_client import AsyncLlamaStackClient`

Applied to files:

  • src/utils/pydantic_ai.py
📚 Learning: 2026-02-25T07:46:39.608Z
Learnt from: asimurka
Repo: lightspeed-core/lightspeed-stack PR: 1211
File: src/models/responses.py:8-16
Timestamp: 2026-02-25T07:46:39.608Z
Learning: In the lightspeed-stack codebase, src/models/requests.py uses OpenAIResponseInputTool as Tool while src/models/responses.py uses OpenAIResponseTool as Tool. This type difference is intentional - input tools and output/response tools have different schemas in llama-stack-api.

Applied to files:

  • src/utils/pydantic_ai.py
📚 Learning: 2026-05-20T08:09:43.391Z
Learnt from: max-svistunov
Repo: lightspeed-core/lightspeed-stack PR: 1580
File: src/llama_stack_configuration.py:651-683
Timestamp: 2026-05-20T08:09:43.391Z
Learning: In `src/llama_stack_configuration.py`, the `apply_high_level_inference` function currently emits `provider_id: p_type` (underscore form, e.g. `sentence_transformers`) directly from the high-level type key, which collides with Llama Stack's hyphenated provider IDs (e.g. `sentence-transformers`). This is a known PoC divergence documented in the spike doc ("Findings discovered during PoC") and tracked in the implementation JIRA "Unified llama_stack.config schema + synthesizer". Decision S5 mandates that each backend-specific synthesizer translates LCORE's canonical type Literal vocabulary to the target backend's expected shape (hyphenated provider_id for Llama Stack; model-string prefixes for Pydantic AI). The PoC code will be removed before merge; the fix belongs in the implementation ticket.

Applied to files:

  • src/utils/pydantic_ai.py
🔇 Additional comments (7)
src/utils/pydantic_ai.py (4)

41-45: 🏗️ Heavy lift

Avoid depending on the protected AsyncLlamaStackClient._client; use the public http_client injection path.

Reaching into client._client is an unstable dependency on a private/internal attribute that can break at runtime even while mocked tests stay green. Prefer constructing/passing the HTTP client via the public injection API.


1-29: LGTM!


48-70: LGTM!


73-104: LGTM!

tests/unit/utils/test_pydantic_ai.py (3)

1-64: LGTM!


66-196: LGTM!


199-270: LGTM!

Comment thread src/utils/pydantic_ai.py
@asimurka
Copy link
Copy Markdown
Contributor

asimurka commented Jun 4, 2026

/retest

@asimurka asimurka requested a review from tisnik June 4, 2026 09:40
Copy link
Copy Markdown
Contributor

@tisnik tisnik left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Comment thread src/utils/pydantic_ai.py
if isinstance(client, AsyncLlamaStackAsLibraryClient):
return LlamaStackProvider(library_client=client)
api_key = client.api_key or "not-needed"
base = str(client.base_url).rstrip("/")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's ok, but IIRC there's some utility function to do it (correctly ;)

@tisnik tisnik merged commit 52b4770 into lightspeed-core:main Jun 4, 2026
43 of 47 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants